#!coding:utf-8
from dis import dis
import cv2
from PIL import Image
import numpy as np
import os
import time

from camera_utils import putText
from pyseeta6 import FaceDetector
from pyseeta6 import FaceAligner
from pyseeta6 import FaceIdentifier

from .ultra_face_inference import UltraFaceInference
from pyseeta6.common import _Face
class FaceRecognizer(object):
    """
    FaceRecognizer类, 用来检测和识别人脸
    Attributes:
        data_dir: str 标记人脸图片的保存目录
        scale: float 缩放倍数, 加快检测速度但会降低检测精度
        font: ImageFont 中文字体
        known_faces: dict 存放已标记人脸的字典
        threshold: float 检测阈值, 小于该值才会进行相似度比较
    """

    def __init__(self, scale=1.0, threshold=0.4, fontSize=18):
        super(FaceRecognizer, self).__init__()
        self.data_dir = os.path.expanduser(
            '~') + "/Lepi_Data/ros/face_recognizer/known_face"
        self.model_dir = os.path.expanduser(
            '~') + "/Lepi_Data/ros/face_recognizer/models"
        # self.font = ImageFont.truetype(
        #     '/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf', fontSize)
        self.known_faces = {}
        self.threshold = threshold
        self.face_locations = []
        self.face_names = []
        self.face_data = []
        self.add_label_success = False
        self.busy = False
        # self.detector = FaceDetector(os.path.join(self.model_dir,'face_detector.csta'),scale=scale)
        self.detector = UltraFaceInference()
        # self.detector.setResize(480, 380)
        self.detector.setResize(240, 180)
        self.detector2 = UltraFaceInference()
        self.detector2.setResize(240, 180)
        self.aligner = FaceAligner(os.path.join(self.model_dir,'face_landmarker_pts5.csta'))
        self.identifier = FaceIdentifier(os.path.join(
            self.model_dir, 'face_recognizer_light.csta'))

        self.load_faces()

    def detect(self, frame):
        """
        人脸检测函数
        Keyword arguments:
        frame: image 原图
        Returns:
        face_locations: [(x1,y1,x2,y2)] 表示所有检测到人脸的位置坐标数组, 每个数组元素代表一个人脸的位置
        """

        if len(frame.shape) == 3:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        start = time.time()
        # faces = self.detector.detect(frame)
        faces = []
        boxes, labels, probs = self.detector.detect(frame)
        for box in boxes:
            face = _Face()
            face.left, face.top, face.right, face.bottom = box
            faces.append(face)
        print('detect %d faces in %d ms' %(len(faces),int(1000*(time.time()-start))))
        self.face_locations = []
        self.face_names = []
        for i, face in enumerate(faces):
            # print('({0},{1},{2},{3}) score={4}'.format(face.left, face.top, face.right, face.bottom, face.score))
            # cv2.rectangle(image_color, (face.left, face.top), (face.right, face.bottom), (0,255,0), thickness=2)
            self.face_locations.append((face.left, face.top, face.right, face.bottom))
            self.face_data = self.getFaceData(0)
            # return self.rect_faces(frame, self.face_locations)
        return faces

    def detect2(self, frame):
        """
        人脸检测函数
        Keyword arguments:
        frame: image 原图
        Returns:
        face_locations: [(x1,y1,x2,y2)] 表示所有检测到人脸的位置坐标数组, 每个数组元素代表一个人脸的位置
        """

        if len(frame.shape) == 3:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        start = time.time()
        # faces = self.detector.detect(frame)
        faces = []
        boxes, labels, probs = self.detector2.detect(frame)
        for box in boxes:
            face = _Face()
            face.left, face.top, face.right, face.bottom = box
            faces.append(face)
        print('detect %d faces in %d ms' %(len(faces),int(1000*(time.time()-start))))
        self.face_locations = []
        self.face_names = []
        for i, face in enumerate(faces):
            # print('({0},{1},{2},{3}) score={4}'.format(face.left, face.top, face.right, face.bottom, face.score))
            # cv2.rectangle(image_color, (face.left, face.top), (face.right, face.bottom), (0,255,0), thickness=2)
            self.face_locations.append((face.left, face.top, face.right, face.bottom))
            self.face_data = self.getFaceData(0)
            # return self.rect_faces(frame, self.face_locations)
        return faces

    def recognize(self, frame):
        """
        人脸识别函数
        Keyword arguments:
        frame: image 原图
        Returns:
        face_locations: [(x1,y1,x2,y2)] 表示所有检测到人脸的位置坐标数组, 每个数组元素代表一个人脸的位置
        face_names: ['未知'] 对应每个人脸的标签, 未检测到用'未知'表示
        """
        if len(frame.shape) == 3:
            image_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        else:
            image_gray = frame
        faces = self.detect2(image_gray)
        if len(faces) == 0:
            return
        # 识别大于10张人脸认为是误检
        if len(faces) > 10:
            self.face_locations = []
            self.face_names = []
            return
        face_names = []
        face_encodings = []
        for face in faces:
            start = time.time()
            if not self.busy:
                self.busy = True
                landmarks = self.aligner.align(image_gray,face)
                feature = self.identifier.extract_feature_with_crop(frame,landmarks)
                self.busy = False
            else:
                aligner = FaceAligner(os.path.join(self.model_dir,'face_landmarker_pts5.csta'))
                identifier = FaceIdentifier(os.path.join(
                    self.model_dir, 'face_recognizer_light.csta'))
                landmarks = aligner.align(image_gray,face)
                feature = identifier.extract_feature_with_crop(frame,landmarks)
                aligner.release()
                identifier.release()
            print('align and extractfeature in %d ms' %(int(1000*(time.time()-start))))
            face_encodings.append(feature)
        # python 3.9 默认生成dict_values 手动转成list进行比较
        known_face_encodings = list(self.known_faces.values())
        known_face_names = list(self.known_faces.keys())
        for face_encoding in face_encodings:
            name = "未知"
            if len(known_face_encodings) > 0:
                # matches = face_recognition.compare_faces(known_face_encodings, face_encoding)
                # print(known_face_encodings, face_encoding)
                face_distances = self.face_distance(
                    known_face_encodings, face_encoding)
                best_match_index = np.argmax(face_distances)
                # print(matches)
                print(face_distances)
                if face_distances[best_match_index] > self.threshold:
                    name = known_face_names[best_match_index]
            face_names.append(name)
        # end = time.time()
        # print("recognized %d faces in %.2f ms" %
        #       (len(face_locations), (end - start)*1000))
        self.face_names = face_names
        self.face_data = self.getFaceData(0)
        # return self.label_faces(frame, self.face_locations, self.face_names)
        # return face_locations, face_names

    def getFaceData(self, index):
        if len(self.face_locations) > index:
            (left, top, right, bottom) = self.face_locations[index]
            return [int((right+left)/2), int((top+bottom)/2), int(right-left), int(bottom-top)]
        else:
            return [0, 0, 0, 0]

    def detectedFaceLabel(self, name):
        if name in self.face_names:
            index = self.face_names.index(name)
            self.face_data = self.getFaceData(index)
            return True
        else:
            self.face_data = []
            return False

    def rect_faces(self, frame, face_locations):
        """
        把检测到的人脸用矩形框出来
        Keyword arguments:
        frame: image 原图
        face_locations: [(x1,y1,x2,y2)] 表示所有检测到人脸的位置坐标数组, 每个数组元素代表一个人脸的位置
        Returns:
        frame: image 框出了人脸的新图像(改变原图)
        """

        # Display the results
        for (left, top, right, bottom) in face_locations:
            # Draw a box around the face
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
            # Draw a label with a name below the face
        return frame

    def label_faces(self, frame, face_locations, face_names):
        """
        给识别到的人脸添加标签
        Keyword arguments:
        frame: image 原图
        face_locations: [(x1,y1,x2,y2)] 表示所有检测到人脸的位置坐标数组, 每个数组元素代表一个人脸的位置
        Returns:
        frame: image 标记了人脸的新图像(不改变原图)
        """
        frame = self.rect_faces(frame, face_locations)

        if len(face_names) == 0:
            return frame

        # Display the results
        for (left, top, right, bottom), name in zip(face_locations, face_names):
            # Draw a box around the face
            # cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
            # Draw a label with a name below the face
            color = (0, 0, 255)
            frame = putText(frame, name, (left + 6, bottom - 24), color)
        return frame

    def load_faces(self):
        """
        加载本地的人脸标签(默认存放在data_dir下)
        Keyword arguments:
        无
        Returns:
        无
        """
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
        files = os.listdir(self.data_dir)
        for file in files:
            try:
                name = file.split('.')[0]
                file_path = os.path.join(self.data_dir, file)
                print(self.add_face_label(cv2.imread(file_path), name))
            except Exception as e:
                print(e)

    def add_face_label(self, frame, name, save=False):
        """
        动态添加人脸标签
        Keyword arguments:
        frame: image 原图
        name: str 标签名称
        save: bool 是否保存, True则将保存图片至本地标签目录, 每次启动会重新读取, False只在本次运行生效
        Returns:
        无
        """

        start = time.time()
        if len(frame.shape) == 3:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        else:
            gray = frame
        faces = self.detect(gray)
        end = time.time()
        print("Found %d faces in %.2f ms" %
              (len(faces), (end - start)*1000))

        if len(faces) == 1:
            start = time.time()
            if not self.busy:
                self.busy = True
                landmarks = self.aligner.align(gray, faces[0])
                feature = self.identifier.extract_feature_with_crop(frame, landmarks)
                self.busy = False
            else:
                aligner = FaceAligner(os.path.join(self.model_dir,'face_landmarker_pts5.csta'))
                identifier = FaceIdentifier(os.path.join(
                    self.model_dir, 'face_recognizer_light.csta'))
                landmarks = aligner.align(gray, faces[0])
                feature = identifier.extract_feature_with_crop(frame, landmarks)
                aligner.release()
                identifier.release()
            self.known_faces[name] = feature
            if save:
                file_path = os.path.join(self.data_dir, name+'.png')
                print(file_path)
                cv2.imwrite(file_path, frame)
            end = time.time()
            print('align and extractfeature in %d ms' %(int(1000*(time.time()-start))))
            
            print('成功添加"%s"的标记' % (name))
            self.add_label_success = True
            return '标记成功'
        elif len(faces) > 1:
            self.add_label_success = False
            return '标记"%s"失败, 检测到多余人脸' % (name)
        else:
            self.add_label_success = False
            return '标记"%s"失败, 未检测到人脸' % (name)

    def remove_face_label(self, name):
        """
        删除人脸标签
        Keyword arguments:
        name: str 标签名称
        Returns:
        无
        """
        if self.known_faces.__contains__(name):
            self.known_faces.pop(name)
        try:
            os.system('rm '+os.path.join(self.data_dir, name+".*"))
            return "已删除"
        except Exception as e:
            print(e)
            return "删除出错"

    def face_distance(self,known_face_encodings, face_encoding):
        dist = []
        for known_face_encoding in known_face_encodings:
            sim = self.identifier.calc_similarity(known_face_encoding, face_encoding)
            dist.append(sim)
        return dist

if __name__ == '__main__':
    pass
